#include <iostream>
#include <vector>
#include <list>

#include <TROOT.h>
#include <TFile.h>
#include <TTree.h>

#include "GeomSvc.h"
#include "TriggerRoad.h"

using namespace std;

int main(int argc, char* argv[])
{
  GeomSvc* p_geomSvc = GeomSvc::instance();
  p_geomSvc->init("geometry_R962");

  Double_t mass, xF, x1, x2;
  Double_t weight;
  Double_t x0, y0, z0;

  Double_t px1, py1, pz1, p1;
  Double_t px2, py2, pz2, p2;

  Int_t nHodoHits1;
  Int_t detectorID1[100];
  Int_t elementID1[100];

  Int_t nHodoHits2;
  Int_t detectorID2[100];
  Int_t elementID2[100];

  TFile* dataFile = new TFile(argv[1], "READ");
  TTree* dataTree = (TTree*)dataFile->Get("save");

  dataTree->SetBranchAddress("weight", &weight);
  dataTree->SetBranchAddress("z0", &z0);
  dataTree->SetBranchAddress("mass", &mass);

  dataTree->SetBranchAddress("nHodoHits1", &nHodoHits1);
  dataTree->SetBranchAddress("detectorID1", detectorID1);
  dataTree->SetBranchAddress("elementID1", elementID1);

  dataTree->SetBranchAddress("px1", &px1);
  dataTree->SetBranchAddress("py1", &py1);

  dataTree->SetBranchAddress("nHodoHits2", &nHodoHits2);
  dataTree->SetBranchAddress("detectorID2", detectorID2);
  dataTree->SetBranchAddress("elementID2", elementID2);

  dataTree->SetBranchAddress("px2", &px2);
  dataTree->SetBranchAddress("py2", &py2);

  TFile* saveFile = new TFile(argv[2], "recreate");
  TTree* saveTree1 = dataTree->CloneTree(0);

  std::list<TriggerRoad> p_roads, m_roads;
  for(int i = 0; i < dataTree->GetEntries(); i++)
    {
      dataTree->GetEntry(i);
      
      /*Add contitions like 
      if(mass < 4) continue;
      if(z0 > 0) continue;
      */

      std::list<TriggerRoad> m_road_add = TriggerRoad::makeRoadList(nHodoHits1, detectorID1, elementID1, z0, mass, sqrt(px1*px1 + py1*py1), weight);
      std::list<TriggerRoad> p_road_add = TriggerRoad::makeRoadList(nHodoHits2, detectorID2, elementID2, z0, mass, sqrt(px2*px2 + py2*py2), weight);
     
      p_roads.insert(p_roads.end(), p_road_add.begin(), p_road_add.end());
      m_roads.insert(m_roads.end(), m_road_add.begin(), m_road_add.end());
      if(p_road_add.empty() || m_road_add.empty()) continue;  

      saveTree1->Fill();
    }

  //Make unique road lists  
  std::list<TriggerRoad> p_roads_unique; p_roads_unique.clear();
  for(std::list<TriggerRoad>::iterator iter = p_roads.begin(); iter != p_roads.end(); ++iter)
    {
      std::list<TriggerRoad>::iterator road = std::find(p_roads_unique.begin(), p_roads_unique.end(), *iter);
      if(road == p_roads_unique.end())
	{
	  p_roads_unique.push_back(*iter);
	}
      else
	{
	  (*road) += (*iter);
	}
    }
  p_roads_unique.sort(TriggerRoad::byTargetDump);

  std::list<TriggerRoad> m_roads_unique; m_roads_unique.clear();
  for(std::list<TriggerRoad>::iterator iter = m_roads.begin(); iter != m_roads.end(); ++iter)
    {
      std::list<TriggerRoad>::iterator road = std::find(m_roads_unique.begin(), m_roads_unique.end(), *iter);
      if(road == m_roads_unique.end())
	{
	  m_roads_unique.push_back(*iter);
	}
      else
	{
	  (*road) += (*iter);
	}
    }
  m_roads_unique.sort(TriggerRoad::byTargetDump);

  ///Assign groupID and unique road ID
  int uniqueID = 0;
  for(std::list<TriggerRoad>::iterator iter = p_roads_unique.begin(); iter != p_roads_unique.end(); ++iter) 
    {
      iter->roadID = uniqueID;
      if(iter->pT_mean > 7.) 
	{
	  iter->groupID = 15;
	}
      else
	{
	  iter->groupID = int(iter->pT_mean/0.5) + 1;
	}
      iter->groupID = iter->groupID*iter->getTB();
      ++uniqueID;
    }

  uniqueID = 0;
  for(std::list<TriggerRoad>::iterator iter = m_roads_unique.begin(); iter != m_roads_unique.end(); ++iter) 
    {
      iter->roadID = uniqueID;
      if(iter->pT_mean > 7.) 
	{
	  iter->groupID = 15;
	}
      else
	{
	  iter->groupID = int(iter->pT_mean/0.5) + 1;
	}
     iter->groupID = iter->groupID*iter->getTB();
     ++uniqueID;
    }

  TriggerRoad* p_road = new TriggerRoad(); p_road->clear();
  TriggerRoad* m_road = new TriggerRoad(); m_road->clear();

  TTree* saveTree2 = new TTree("single_p", "single_p");
  saveTree2->Branch("road", &p_road, 256000, 99);
  for(std::list<TriggerRoad>::iterator iter = p_roads_unique.begin(); iter != p_roads_unique.end(); ++iter)
    {
      *p_road = *iter;
      saveTree2->Fill();
    }

  TTree* saveTree3 = new TTree("single_m", "single_m");
  saveTree3->Branch("road", &m_road, 256000, 99);
  for(std::list<TriggerRoad>::iterator iter = m_roads_unique.begin(); iter != m_roads_unique.end(); ++iter)
    {
      *m_road = *iter;
      saveTree3->Fill();
    }

  saveFile->cd();
  saveTree1->Write();
  saveTree2->Write();
  saveTree3->Write();
  saveFile->Close();

  return 1;
}
